#!/usr/bin/env python  
#-*- coding:utf-8 _*-  
""" 
@author:hello_life 
@license: Apache Licence 
@file: bilstm_crf.py 
@time: 2022/05/01
@software: PyCharm 
description:
"""
import torch
import torch.nn as nn
from torchcrf import CRF
import numpy as np

class Bilstm_CRF(nn.Module):

    def __init__(self, config):
        super(Bilstm_CRF,self).__init__()
        self.embedding_pretrained=torch.tensor(np.load(config.pretrain_save_dir)["embedding"],dtype=torch.float)
        self.embedding=nn.Embedding.from_pretrained(self.embedding_pretrained,freeze=False)
        self.lstm=nn.LSTM(300,config.hidden_size,2,bidirectional=True,batch_first=True)
        self.l1=nn.Linear(config.hidden_size*2,config.num_class)
        self.crf=CRF(config.num_class,batch_first=True)

    def forward(self,x):
        emb=self.embedding(x)
        output,(h_0,c_0)=self.lstm(emb)
        tag_scores=self.l1(output)
        return tag_scores

    def forward_with_crf(self,x,mask,y):
        tag_scores=self.forward(x)
        loss=self.crf(tag_scores,y,mask)*(-1)
        return tag_scores,loss



